# 一段python代码，可以将.csv的数据集分成训练集，验证集和测试集，比例大小自行设置
import pandas as pd
from sklearn.model_selection import train_test_split
from pathlib import Path


def split_csv(csv_file_path,
              train_pro, val_pro, test_pro,
              csv_encoding,
              out_to_src_folder=True,
              train_out='train.csv', val_out='val.csv', test_out='test.csv'):
    #载入数据
    data = pd.read_csv(csv_file_path, encoding=csv_encoding)

    #将数据集分为训练集，验证集和测试集
    train_val_data, test_data=train_test_split(data, test_size=test_pro/(train_pro+val_pro+test_pro), random_state=42)
    train_data, val_data=train_test_split(train_val_data, test_size=val_pro/(train_pro+val_pro), random_state=42)

    # 获取csv文件所在的目录
    dir_path = Path(csv_file_path).parent

    #保存训练集，验证集和测试集
    if out_to_src_folder:
        train_data.to_csv(dir_path/train_out, index=False)
        val_data.to_csv(dir_path/val_out, index=False)
        test_data.to_csv(dir_path/test_out, index=False)
        print('保存路径是：', dir_path)
    else:
        train_data.to_csv(train_out, index=False)
        val_data.to_csv(val_out, index=False)
        test_data.to_csv(test_out, index=False)

if __name__ == "__main__":
    csv_file_path = r'C:\MyNuts\kerasStudy\paper01_new\data(exportFromRasterLayers)\RasterUnits20240422\研究区_pt.csv'
    # 训练集，验证集和测试集比例为6:2:2
    train_pro, val_pro, test_pro = 6, 2, 2
    out_to_src_folder = True
    train_out = 'train.csv'
    val_out = 'val.csv'
    test_out = 'test.csv'
    split_csv(csv_file_path,
              train_pro, val_pro, test_pro,
              'utf-8',
              out_to_src_folder,
              train_out, val_out, test_out)
